package com.study.mongodbdynamicdatasource.config;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.lang.Assert;
import cn.hutool.core.util.ReflectUtil;
import cn.hutool.core.util.StrUtil;
import com.mongodb.ClientSessionOptions;
import com.mongodb.client.MongoCollection;
import com.mongodb.client.MongoDatabase;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.bson.Document;
import org.springframework.data.geo.GeoResults;
import org.springframework.data.mongodb.MongoDatabaseFactory;
import org.springframework.data.mongodb.core.*;
import org.springframework.data.mongodb.core.aggregation.Aggregation;
import org.springframework.data.mongodb.core.aggregation.AggregationOperationContext;
import org.springframework.data.mongodb.core.aggregation.AggregationResults;
import org.springframework.data.mongodb.core.convert.MongoWriter;
import org.springframework.data.mongodb.core.mapreduce.GroupBy;
import org.springframework.data.mongodb.core.mapreduce.GroupByResults;
import org.springframework.data.mongodb.core.mapreduce.MapReduceOptions;
import org.springframework.data.mongodb.core.mapreduce.MapReduceResults;
import org.springframework.data.mongodb.core.query.*;
import org.springframework.lang.Nullable;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * @Name: MultiMongoTemplate
 * @Description: java类作用描述
 * @Author: yn
 * @CreateDate: 2023/6/9 15:31
 * @Remark: 特殊说明
 * @Version: 1.0
 */
@Slf4j
public class MultiMongoTemplate extends MongoTemplate {

	public final IgnoreCollectionProperties IGNORE_COLLECTION_PROPERTIES;

	public MultiMongoTemplate(IgnoreCollectionProperties ignoreCollectionProperties, MongoDatabaseFactory mongoDbFactory) throws Exception {
		super(mongoDbFactory);
		MongoDbContext.setDefaultMongoDatabaseFactory(mongoDbFactory);
		IGNORE_COLLECTION_PROPERTIES=ignoreCollectionProperties;
	}

	@Override
	protected MongoDatabase doGetDatabase() {
		MongoDatabaseFactory mongoDbFactory = MongoDbContext.getMongoDbFactory();
		return mongoDbFactory == null ? super.doGetDatabase() : mongoDbFactory.getMongoDatabase();
	}

	@Override
	public SessionScoped withSession(ClientSessionOptions options) {
		Assert.notNull(options, "ClientSessionOptions must not be null!");
		MongoDatabaseFactory mongoDbFactory = MongoDbContext.getMongoDbFactory()==null?getMongoDbFactory():MongoDbContext.getMongoDbFactory();
		return withSession(() -> mongoDbFactory.getSession(options));
	}

	private void switchDbFactory(String tenantId) {
		MongoDbContext.pushMongoDatabaseFactory(tenantId);
	}

	private void checkAndSwitchDbFactory(String collectionName) {
		if (checkCollection(collectionName)) {
			switchDbFactory(AuthUtil.getTenantId());
		}
	}

	private boolean checkCollection(String collectionName) {
		if (CollectionUtil.isNotEmpty(IGNORE_COLLECTION_PROPERTIES.getCollections()) && IGNORE_COLLECTION_PROPERTIES.getCollections().contains(collectionName)) {
			return false;
		}
		return true;
	}

	private void checkAndSwitchDbFactoryByQuery(Query query, String collectionName) {
		checkAndSwitchDbFactoryByDocument(query.getFieldsObject(),collectionName);
	}

	private void checkAndSwitchDbFactoryByDocument(Document query, String collectionName) {
		if (checkCollection(collectionName)) {
			if (!query.containsKey("tenantId")) {
				query.put("tenantId",AuthUtil.getTenantId());
			}
			switchDbFactory(query.get("tenantId", String.class));
		}
	}

	@SneakyThrows
	private <T> void checkAndSwitchDbFactoryBySaveEntity(T objectToSave, String collectionName) {
		if (checkCollection(collectionName)) {
			Boolean hasTenantId = ReflectUtil.hasField(objectToSave.getClass(), "tenantId");
			if (hasTenantId) {
				String tenantId = (String) ReflectUtil.getFieldValue(objectToSave, "tenantId");
				if (StrUtil.isBlank(tenantId)) {
					ReflectUtil.setFieldValue(objectToSave, "tenantId",AuthUtil.getTenantId());
					tenantId=AuthUtil.getTenantId();
				}
				switchDbFactory(tenantId);
			}
		}
	}

	private void clearDbContext() {
	}

	////////////////////
	@Override
	protected void executeQuery(Query query, String collectionName, DocumentCallbackHandler documentCallbackHandler, CursorPreparer preparer) {
		checkAndSwitchDbFactoryByQuery(query,collectionName);
		super.executeQuery(query, collectionName, documentCallbackHandler, preparer);
		clearDbContext();
	}

	//execute 不重写 基本也不会调用
//	@Override
//	public <T> T execute(String collectionName, CollectionCallback<T> callback) {
//		checkAndSwitchDbFactory(collectionName);
//		T x = super.execute(collectionName, callback);
//		clearDbContext();
//		return x;
//	}

	@Override
	public MongoCollection<Document> getCollection(String collectionName) {
		checkAndSwitchDbFactory(collectionName);
		MongoCollection<Document> x = super.getCollection(collectionName);
		clearDbContext();
		return x;
	}

	@Override
	public boolean collectionExists(String collectionName) {
		checkAndSwitchDbFactory(collectionName);
		boolean x = super.collectionExists(collectionName);
		clearDbContext();
		return x;
	}

	@Override
	public <T> List<T> findDistinct(Query query, String field, String collectionName, Class<?> entityClass, Class<T> resultClass) {
		checkAndSwitchDbFactoryByQuery(query,collectionName);
		List<T> x = super.findDistinct(query, field, collectionName, entityClass, resultClass);
		clearDbContext();
		return x;
	}

	@Override
	public <T> GeoResults<T> geoNear(NearQuery near, Class<?> domainType, String collectionName, Class<T> returnType) {
		checkAndSwitchDbFactory(collectionName);
		GeoResults<T> x = super.geoNear(near, domainType, collectionName, returnType);
		clearDbContext();
		return x;
	}

	@Override
	protected <T> T doInsert(String collectionName, T objectToSave, MongoWriter<T> writer) {
		checkAndSwitchDbFactoryBySaveEntity(objectToSave,collectionName);
		T x = super.doInsert(collectionName, objectToSave, writer);
		clearDbContext();
		return x;
	}

	@Override
	protected <T> Collection<T> doInsertBatch(String collectionName, Collection<? extends T> batchToSave, MongoWriter<T> writer) {
		if (checkCollection(collectionName)) {
			if (CollectionUtil.isEmpty(batchToSave)) {
				return (Collection<T>) batchToSave;
			}
			Map<String, ? extends List<? extends T>> tenantIdGroup = batchToSave.stream().collect(Collectors.groupingBy((i -> (String) ReflectUtil.getFieldValue(i, "tenantId"))));
			Collection<T> x =new ArrayList<>();
			tenantIdGroup.forEach((k,v)->{
				switchDbFactory(k);
				x.addAll(super.doInsertBatch(collectionName, v, writer));
			});
			clearDbContext();
			return x;
		}else{
			Collection<T> x = super.doInsertBatch(collectionName, batchToSave, writer);
			return x;
		}
	}

	@Override
	protected <T> T doSave(String collectionName, T objectToSave, MongoWriter<T> writer) {
		checkAndSwitchDbFactoryBySaveEntity(objectToSave, collectionName);
		T x = super.doSave(collectionName, objectToSave, writer);
		clearDbContext();
		return x;
	}

	@Override
	public <T> List<T> findAll(Class<T> entityClass, String collectionName) {
		checkAndSwitchDbFactory(collectionName);
		List<T> x = super.findAll(entityClass, collectionName);
		clearDbContext();
		return x;
	}

	@Override
	public <T> MapReduceResults<T> mapReduce(Query query, String inputCollectionName, String mapFunction, String reduceFunction, MapReduceOptions mapReduceOptions, Class<T> entityClass) {
		checkAndSwitchDbFactoryByQuery(query,inputCollectionName);
		MapReduceResults<T> x = super.mapReduce(query, inputCollectionName, mapFunction, reduceFunction, mapReduceOptions, entityClass);
		clearDbContext();
		return x;
	}

	@Override
	public <T> List<T> mapReduce(Query query, Class<?> domainType, String inputCollectionName, String mapFunction, String reduceFunction, MapReduceOptions mapReduceOptions, Class<T> resultType) {
		checkAndSwitchDbFactoryByQuery(query,inputCollectionName);
		List<T> x = super.mapReduce(query, domainType, inputCollectionName, mapFunction, reduceFunction, mapReduceOptions, resultType);
		clearDbContext();
		return x;
	}

	@Override
	public <T> GroupByResults<T> group(Criteria criteria, String inputCollectionName, GroupBy groupBy, Class<T> entityClass) {
		checkAndSwitchDbFactory(inputCollectionName);
		GroupByResults<T> x = super.group(criteria, inputCollectionName, groupBy, entityClass);
		clearDbContext();
		return x;
	}

	@Override
	protected <O> AggregationResults<O> doAggregate(Aggregation aggregation, String collectionName, Class<O> outputType, AggregationOperationContext context) {
		checkAndSwitchDbFactory(collectionName);
		AggregationResults<O> x = super.doAggregate(aggregation, collectionName, outputType, context);
		clearDbContext();
		return x;
	}

	@Override
	protected MongoCollection<Document> doCreateCollection(String collectionName, Document collectionOptions) {
		checkAndSwitchDbFactory(collectionName);
		MongoCollection<Document> x = super.doCreateCollection(collectionName, collectionOptions);
		clearDbContext();
		return x;
	}

	@Override
	protected <T> T doFindOne(String collectionName, Document query, Document fields, Class<T> entityClass) {
		checkAndSwitchDbFactoryByDocument(query, collectionName);
		T x = super.doFindOne(collectionName, query, fields, entityClass);
		clearDbContext();
		return x;
	}

	@Override
	protected <T> List<T> doFind(String collectionName, Document query, Document fields, Class<T> entityClass,
								 CursorPreparer preparer) {
		checkAndSwitchDbFactoryByDocument(query, collectionName);
		List<T> x = super.doFind(collectionName, query, fields, entityClass, preparer);
		clearDbContext();
		return x;
	}

	@Override
	protected <T> List<T> doFind(String collectionName, Document query, Document fields, Class<T> entityClass) {
		checkAndSwitchDbFactoryByDocument(query, collectionName);
		List<T> x = super.doFind(collectionName, query, fields, entityClass);
		clearDbContext();
		return x;
	}

	@Override
	protected <T> T doFindAndRemove(String collectionName, Document query, Document fields, Document sort, Collation collation, Class<T> entityClass) {
		checkAndSwitchDbFactoryByDocument(query, collectionName);
		T x = super.doFindAndRemove(collectionName, query, fields, sort, collation, entityClass);
		clearDbContext();
		return x;
	}

	@Override
	protected <T> T doFindAndModify(String collectionName, Document query, Document fields, Document sort,
									Class<T> entityClass, UpdateDefinition update, @Nullable FindAndModifyOptions options) {
		checkAndSwitchDbFactoryByDocument(query, collectionName);
		T x = super.doFindAndModify(collectionName, query, fields, sort, entityClass, update, options);
		clearDbContext();
		return x;
	}

	@Override
	protected <T> T doFindAndReplace(String collectionName, Document mappedQuery, Document mappedFields, Document mappedSort, com.mongodb.client.model.Collation collation, Class<?> entityType, Document replacement, FindAndReplaceOptions options, Class<T> resultType) {
		checkAndSwitchDbFactoryByDocument(mappedQuery, collectionName);
		T x = super.doFindAndReplace(collectionName, mappedQuery, mappedFields, mappedSort, collation, entityType, replacement, options, resultType);
		clearDbContext();
		return x;
	}
}
